# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Molecule Embedding
"""

from typing import Union, Tuple, Optional

import torch
from torch import nn, Tensor

from ..units import GLOBAL_UNITS, Length, get_length
from ..utils import get_integer, get_tensor, get_initializer, get_arguments
from ..cutoff import Cutoff, get_cutoff
from ..rbf import get_rbf
from ..filter import Filter, get_filter
from .graph import GraphEmbedding, _embedding_register


@_embedding_register('molecule')
class MolEmbedding(GraphEmbedding):
    """
    Molecule embedding module.
    
    Args:
        dim_node (int): Dimension of node embeddings.
        dim_edge (int, optional): Dimension of edge embeddings. Defaults to None.
        emb_dis (bool, optional): Whether to embed distance information. Defaults to True.
        emb_bond (bool, optional): Whether to embed bond information. Defaults to False.
        cutoff (Length): Cutoff distance for interactions. Defaults to Length(1, 'nm').
        cutoff_fn (str, optional): Type of cutoff function. Defaults to 'smooth'.
        rbf_fn (str, optional): Type of radial basis function. Defaults to 'log_gaussian'.
        num_basis (int, optional): Number of basis functions. Defaults to None.
        atom_filter (str, optional): Type of atom filter. Defaults to None.
        dis_filter (str, optional): Type of distance filter. Defaults to None.
        bond_filter (str, optional): Type of bond filter. Defaults to 'residual'.
        interaction (callable, optional): Interaction function for combining embeddings.
        dis_self (Length, optional): Self-distance for atoms. Defaults to Length(0.05, 'nm').
        use_sub_cutoff (bool, optional): Whether to use sub-cutoff. Defaults to False.
        cutoff_buffer (Length, optional): Buffer for cutoff distance. Defaults to Length(0.2, 'nm').
        num_atom_types (int, optional): Number of atom types. Defaults to 64.
        num_bond_types (int, optional): Number of bond types. Defaults to 16.
        initializer (callable, optional): Weight initializer. Defaults to None.
        activation (str, optional): Activation function. Defaults to 'silu'.
        length_unit (str, optional): Length unit. Defaults to GLOBAL_UNITS.length_unit.
    """
    def __init__(self, 
                 dim_node: int, 
                 dim_edge: Optional[int] = None, 
                 emb_dis: bool = True, 
                 emb_bond: bool = False, 
                 cutoff: Length = Length(1, 'nm'),
                 cutoff_fn: str = 'smooth', 
                 rbf_fn: str = 'log_gaussian', 
                 num_basis: Optional[int] = None, 
                 atom_filter: Optional[str] = None, 
                 dis_filter: Optional[str] = None,
                 bond_filter: str = 'residual', 
                 interaction: Optional[callable] = None, 
                 dis_self: Length = Length(0.05, 'nm'), 
                 use_sub_cutoff: bool = False,
                 cutoff_buffer: Length = Length(0.2, 'nm'), 
                 num_atom_types: int = 64, 
                 num_bond_types: int = 16, 
                 initializer: Optional[callable] = None,
                 activation: str = 'silu', 
                 length_unit: str = GLOBAL_UNITS.length_unit,
                 **kwargs):
        
        super().__init__(
            dim_node=dim_node, 
            dim_edge=dim_node if dim_edge is None else dim_edge,
            emb_dis=emb_dis, 
            emb_bond=emb_bond, 
            activation=activation,
            length_unit=length_unit
        )
        self._kwargs = get_arguments(locals(), kwargs)

        self.num_atom_types = num_atom_types
        self.atom_embedding = nn.Embedding(
            self.num_atom_types,
            self.dim_node,
            device=self.device
            )
        
        self.cutoff = get_length(cutoff, self.units)
        self.cutoff_fn = get_cutoff(cutoff_fn, self.cutoff)
        if self.cutoff_fn is not None:
            self.cutoff = self.cutoff_fn.cutoff
        if self.cutoff is not None:
            self.cutoff = get_tensor(self.cutoff, dtype=torch.float32, device=self.device)
        
        dis_self = get_length(dis_self, self.units)
        self.dis_self = get_tensor(dis_self, dtype=torch.float32, device=self.device).reshape((-1,))
        self.num_bond_types = num_bond_types

        self.atom_filter = get_filter(
            atom_filter,
            self.dim_node,
            self.dim_node,
            activation
            )
        
        self.use_sub_cutoff = use_sub_cutoff
        self.cutoff_buffer = get_length(cutoff_buffer, self.units)

        self.rbf_fn = None
        self.dis_filter = None
        if emb_dis:
            num_basis = get_integer(num_basis)
            self.rbf_fn = get_rbf(
                rbf_fn,
                r_max=self.cutoff,
                num_basis=num_basis,
                length_unit=self.units.length_unit
                )
            self.dis_filter = get_filter(
                cls_name=dis_filter,
                dim_in=self.num_basis,
                dim_out=self.dim_edge,
                activation=activation
                )
            
        self.bond_embedding = None
        self.bond_filter = None
        if emb_bond:
            self.bond_embedding = nn.Embedding(
                self.num_bond_types,
                self.dim_edge,
                device=self.device
                )
            self.bond_filter = get_filter(
                cls_name=bond_filter,
                dim_in=self.dim_edge,
                dim_out=self.dim_edge,
                activation=activation
                )
        if emb_dis and emb_bond:
            self.interaction = interaction

    @property
    def num_basis(self) -> int:
        """number of radical basis function"""
        if self.rbf_fn is None:
            return 1
        return self.rbf_fn.num_basis

    @property
    def dim_edge(self) -> int:
        """dimension of edge vector"""
        if self.emb_dis and self.dis_filter is None:
            return self.num_basis
        return self._dim_edge
    
    def _print_info(self,
                    num_retraction: int = 3,
                    num_gap: int = 3,
                    char: str = ' '
                    ) -> None:
        """Print information about the embedding module."""
        ret = char * num_retraction
        gap = char * num_gap
        print(ret + f' Graph Embedding: {self._get_name()}')
        print('-' * 80)
        print(ret + gap + f' Length unit: {self.units.length_unit}')
        print(ret + gap + f' Atom embedding size: {self.num_atom_types}')
        print(ret + gap + f' Cutoff distance: {self.cutoff} {self.units.length_unit}')
        print(ret + gap + f' Cutoff function: {self.cutoff_fn._get_name()}')
        print(ret + gap + f' Radical basis functions: {self.rbf_fn._get_name()}')
        self.rbf_fn._print_info(
            num_retraction=num_retraction + num_gap, 
            num_gap=num_gap, 
            char=char
        )
        print(ret + gap + f' Embedding distance: {self.emb_dis}')
        print(ret + gap + f' Embedding Bond: {self.emb_bond}')
        print(ret + gap + f' Dimension of node embedding vector: {self.dim_node}')
        print(ret + gap + f' Dimension of edge embedding vector: {self.dim_edge}')
        print('-' * 80)

    def get_rbf(self, distances: Tensor) -> Tensor:
        """get radical basis function"""
        if self.rbf_fn is None:
            # (B, A, N, 1)
            return distances.unsqueeze(-1)
        # (B, A, N, F)
        return self.rbf_fn(distances)

    def forward(self, 
                atom_type: Tensor, 
                atom_mask: Tensor, 
                distance: Tensor, 
                dis_mask: Tensor, 
                bond: Optional[Tensor] = None, 
                bond_mask: Optional[Tensor] = None, 
                neigh_idx: Optional[Tensor] = None
                ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        """   
        Args:
            atom_type: Tensor of atom types.
            atom_mask: Tensor indicating valid atoms.
            distance: Tensor of distances between atoms.
            dis_mask: Tensor indicating valid distances.
            bond: Optional tensor of bond types.
            bond_mask: Optional tensor indicating valid bonds.
            neigh_idx: Optional tensor of neighbor indices.
        
        Returns:
            node_emb: Tensor of node embeddings.
            node_mask: Tensor indicating valid nodes.
            edge_emb: Tensor of edge embeddings.
            edge_mask: Tensor indicating valid edges.
            edge_cutoff: Tensor of edge cutoffs.
        """
        if self.emb_dis:
            batch_size = distance.shape[0]
            num_atoms = distance.shape[-2]
        else:
            batch_size = bond.shape[0]
            num_atoms = bond.shape[-2]
        
        node_emb = self.atom_embedding(atom_type)
        if self.atom_filter is not None:
            node_emb = self.atom_filter(node_emb)

        node_mask = atom_mask
        if batch_size > 1 and atom_type.shape[0] != batch_size:
            node_emb = node_emb.broadcast_to((batch_size,) + node_emb.shape[1:])
            if atom_mask is not None:
                node_mask = node_mask.broadcast_to((batch_size,) + node_mask.shape[1:])
        
        dis_emb = None
        _dis_mask = None
        dis_cutoff = None
        if self.emb_dis:
            if neigh_idx is None:
                distance = torch.where(
                    torch.eye(num_atoms,num_atoms,dtype=torch.bool, device=distance.device),
                    self.dis_self.to(distance.device),
                    distance)
            else:
                distance[:,:,0] = self.dis_self.to(distance.device)

            # Calculate RBF and apply filter
            dis_emb = self.get_rbf(distance)
            if self.dis_filter is not None:
                dis_emb = self.dis_filter(dis_emb)

            # Calculate RBF and apply filter
            if self.cutoff_fn is None:
                dis_cutoff = torch.ones_like(distance, device=self.device)
            else:
                if self.use_sub_cutoff:
                    center_dis = torch.unsqueeze(distance[..., 0, :], -2)
                    cutoff = self.cutoff + self.cutoff_buffer - center_dis
                    cutoff = torch.maximum(torch.tensor([1e-8],device=cutoff.device), torch.minimum(cutoff, self.cutoff))
                    dis_cutoff, _dis_mask = self.cutoff_fn(distance, dis_mask, cutoff)
                else:
                    dis_cutoff, _dis_mask = self.cutoff_fn(distance, dis_mask)
        
        if dis_mask is None:
            dis_mask = _dis_mask
        else:
            dis_mask = torch.logical_and(dis_mask, _dis_mask)
        
        bond_emb = None
        bond_cutoff = None
        if self.emb_bond:
            bond_emb = self.bond_embedding(bond)
            if bond_mask is not None:
                bond_emb = bond_emb * bond_mask.unsqueeze(-1)
                bond_cutoff = (bond_mask > 0).to(bond_emb.dtype)
            if self.bond_filter is not None:
                bond_emb = self.bond_filter(bond_emb)
        
        edge_cutoff = dis_cutoff
        edge_mask = dis_mask
        if not self.emb_dis:
            edge_emb = bond_emb
            edge_mask = bond_mask
            edge_cutoff = bond_cutoff
        elif not self.emb_bond:
            edge_emb = dis_emb
        else:
            node_emb, edge_emb = self.interaction(
                node_emb,
                node_emb,
                bond_emb,
                bond_mask,
                bond_cutoff
                )
        return node_emb, node_mask, edge_emb, edge_mask, edge_cutoff
